{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "from eval_utils import *\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "import BboxToolkit as bt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "import inflect\n",
    "\n",
    "# Create an engine instance\n",
    "convert = inflect.engine()\n",
    "\n",
    "def find_numbers(text):\n",
    "    # Define a pattern to match numerical digits\n",
    "    number_pattern = r\"\\b\\d+\\b\"\n",
    "\n",
    "    # Define a pattern to match number words (adjust according to needed range or specifics)\n",
    "    word_pattern = r\"\\b(one|two|three|four|five|six|seven|eight|nine|ten)\\b\"\n",
    "\n",
    "    # Find all occurrences of numerical digits\n",
    "    numbers = re.findall(number_pattern, text)\n",
    "\n",
    "    # Find all occurrences of number words\n",
    "    number_words = re.findall(word_pattern, text, flags=re.IGNORECASE)\n",
    "\n",
    "    # You might want to convert found words to their corresponding numbers\n",
    "    word_to_number = {\n",
    "        \"one\": 1, \"two\": 2, \"three\": 3, \"four\": 4, \"five\": 5,\n",
    "        \"six\": 6, \"seven\": 7, \"eight\": 8, \"nine\": 9, \"ten\": 10\n",
    "    }\n",
    "    numbers.extend(str(word_to_number[word.lower()]) for word in number_words)\n",
    "\n",
    "    return numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import shapely.geometry as shgeo\n",
    "\n",
    "from BboxToolkit import bbox2type\n",
    "from BboxToolkit import get_bbox_type\n",
    "\n",
    "def bbox_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-6):\n",
    "    assert mode in ['iou', 'iof']\n",
    "    assert get_bbox_type(bboxes1) != 'notype'\n",
    "    assert get_bbox_type(bboxes2) != 'notype'\n",
    "    rows = bboxes1.shape[0]\n",
    "    cols = bboxes2.shape[0]\n",
    "    if is_aligned:\n",
    "        assert rows == cols\n",
    "\n",
    "    if rows * cols == 0:\n",
    "        return np.zeros((rows, 1), dtype=np.float32) \\\n",
    "                if is_aligned else np.zeros((rows, cols), dtype=np.float32)\n",
    "\n",
    "    hbboxes1 = bbox2type(bboxes1, 'hbb')\n",
    "    hbboxes2 = bbox2type(bboxes2, 'hbb')\n",
    "    if not is_aligned:\n",
    "        hbboxes1 = hbboxes1[:, None, :]\n",
    "    lt = np.maximum(hbboxes1[..., :2], hbboxes2[..., :2])\n",
    "    rb = np.minimum(hbboxes1[..., 2:], hbboxes2[..., 2:])\n",
    "    wh = np.clip(rb - lt, 0, np.inf)\n",
    "    h_overlaps = wh[..., 0] * wh[..., 1]\n",
    "\n",
    "    if get_bbox_type(bboxes1) == 'hbb' and get_bbox_type(bboxes2) == 'hbb':\n",
    "        overlaps = h_overlaps\n",
    "        areas1 = (hbboxes1[..., 2] - hbboxes1[..., 0]) * (\n",
    "            hbboxes1[..., 3] - hbboxes1[..., 1])\n",
    "\n",
    "        if mode == 'iou':\n",
    "            areas2 = (hbboxes2[..., 2] - hbboxes2[..., 0]) * (\n",
    "                hbboxes2[..., 3] - hbboxes2[..., 1])\n",
    "            unions = areas1 + areas2 - overlaps\n",
    "        else:\n",
    "            unions = areas1\n",
    "\n",
    "\n",
    "    else:\n",
    "        polys1 = bbox2type(bboxes1, 'poly')\n",
    "        polys2 = bbox2type(bboxes2, 'poly')\n",
    "        sg_polys1 = [shgeo.Polygon(p) for p in polys1.reshape(rows, -1, 2)]\n",
    "        sg_polys2 = [shgeo.Polygon(p) for p in polys2.reshape(cols, -1, 2)]\n",
    "\n",
    "        overlaps = np.zeros(h_overlaps.shape)\n",
    "        for p in zip(*np.nonzero(h_overlaps)):\n",
    "            overlaps[p] = sg_polys1[p[0]].intersection(sg_polys2[p[-1]]).area\n",
    "\n",
    "        if mode == 'iou':\n",
    "            unions = np.zeros(h_overlaps.shape, dtype=np.float32)\n",
    "            for p in zip(*np.nonzero(h_overlaps)):\n",
    "                unions[p] = sg_polys1[p[0]].union(sg_polys2[p[-1]]).area\n",
    "        else:\n",
    "            unions = np.array([p.area for p in sg_polys1], dtype=np.float32)\n",
    "            if not is_aligned:\n",
    "                unions = unions[..., None]\n",
    "\n",
    "    unions = np.clip(unions, eps, np.inf)\n",
    "    outputs = overlaps / unions\n",
    "    if outputs.ndim == 1:\n",
    "        outputs = outputs[..., None]\n",
    "    return outputs, overlaps, unions\n",
    "\n",
    "def corner_to_obb(corners):\n",
    "    # Extract x1, y1, x2, y2, angle from the corners array\n",
    "    x1 = corners[:, 0]\n",
    "    y1 = corners[:, 1]\n",
    "    x2 = corners[:, 2]\n",
    "    y2 = corners[:, 3]\n",
    "    angle = corners[:, 4]\n",
    "    angle = angle % 180\n",
    "    \n",
    "    theta_rad = np.radians(angle)\n",
    "\n",
    "    # Calculate the midpoint (center of the bounding box)\n",
    "    cx = (x1 + x2) / 2.0\n",
    "    cy = (y1 + y2) / 2.0\n",
    "\n",
    "    # Calculate distance (width and height) between points\n",
    "    # Width is along the major axis after rotation\n",
    "    dx = x2 - x1\n",
    "    dy = y2 - y1\n",
    "    w = np.abs(np.sqrt(dx**2 + dy**2) * np.cos(theta_rad))\n",
    "    h = np.abs(np.sqrt(dx**2 + dy**2) * np.sin(theta_rad))\n",
    "\n",
    "    # if w<h: # [45,135]\n",
    "    #     w,h = h,w\n",
    "    #     angle = 90 - angle\n",
    "\n",
    "    w, h = np.maximum(w, h), np.minimum(w, h)\n",
    "    \n",
    "    # Combine results into a single 2D array\n",
    "    angle = angle % 180\n",
    "    results = np.vstack((cx, cy, w, h, angle)).T\n",
    "\n",
    "    return results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "output_dir = '../../outputs/outputs_RSBench_new/'\n",
    "# output_dir = '../../outputs/outputs_RSBench_obb_v2_new/'\n",
    "\n",
    "# output_version =  '_v1' \n",
    "output_version =  '_v2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path= os.path.join(output_dir, f'rsbench_referring{output_version}.json')\n",
    "# data_path= os.path.join(output_dir, f'rsvgd.json')\n",
    "\n",
    "thres_list = [0.5, 0.7]\n",
    "count = np.zeros(len(thres_list))\n",
    "\n",
    "cumI = 0\n",
    "cumU = 0\n",
    "mean_IoU = 0\n",
    "\n",
    "total_count = 0\n",
    "valid_count = 0\n",
    "\n",
    "use_size_group = False\n",
    "size_list = ['small', 'medium', 'large']\n",
    "unique_check = [1,0]\n",
    "rot_errs = []\n",
    "\n",
    "with open(data_path, 'r') as file:\n",
    "    for line in file:\n",
    "        # Convert JSON string to Python dictionary\n",
    "        item = json.loads(line.strip())\n",
    "        img_id = item['image_id']\n",
    "        is_unique =  item['is_unique']\n",
    "        if use_size_group:\n",
    "            obj_size = item['size_group']\n",
    "            obj_size = 'medium' if obj_size == '' else obj_size\n",
    "        else:\n",
    "            obj_size = 'medium'\n",
    "\n",
    "        if not (is_unique in unique_check and (obj_size in size_list)):\n",
    "            continue\n",
    "        \n",
    "        total_count += 1\n",
    "\n",
    "        if 'obb_v1' in output_dir: \n",
    "            integers = re.findall(r'\\d+', item['obb_v1'])\n",
    "        elif 'obb_v2' in output_dir:\n",
    "            integers = re.findall(r'\\d+', item['obb_v2'])\n",
    "        else:\n",
    "            integers = re.findall(r'\\d+', item['ground_truth'])\n",
    "        gt_bbox = np.array([int(num) for num in integers])[np.newaxis,:]\n",
    "        \n",
    "        integers = re.findall(r'\\d+', item['answer'])\n",
    "        pred_bbox = np.array([int(num) for num in integers])[np.newaxis,:]\n",
    "        \n",
    "        try:\n",
    "            if 'obb_v1' in output_dir and len(integers) > 4:\n",
    "                rot_errs.append(pred_bbox[0,4] % 180 - gt_bbox[0,4] % 180)\n",
    "                iou_score, I, U = bbox_overlaps(gt_bbox, pred_bbox)\n",
    "            elif 'obb_v2' in output_dir and len(integers) > 4:\n",
    "                gt_bbox = corner_to_obb(gt_bbox)\n",
    "                pred_bbox = corner_to_obb(pred_bbox)\n",
    "                rot_errs.append(pred_bbox[0,4] - gt_bbox[0,4])\n",
    "                iou_score, I, U = bbox_overlaps(gt_bbox, pred_bbox)\n",
    "            elif 'obb' not in output_dir:\n",
    "                # gt_bbox = gt_bbox[:,:4]\n",
    "                # pred_bbox = pred_bbox[:,:4]\n",
    "                # iou_score, I, U = bbox_overlaps(gt_bbox, pred_bbox)\n",
    "                gt_bbox = gt_bbox[0,:4].tolist()\n",
    "                pred_bbox = pred_bbox[0,:4].tolist()\n",
    "                iou_score, I, U = computeIoU(gt_bbox, pred_bbox, return_iou=True)\n",
    "                iou_score = [iou_score]\n",
    "                I = [I]\n",
    "                U = [U]\n",
    "            else:\n",
    "                print('invalid output', img_id, pred_bbox, gt_bbox, flush=True)\n",
    "                continue\n",
    "\n",
    "            iou_score = iou_score[0]\n",
    "            I = I[0]\n",
    "            U = U[0]\n",
    "\n",
    "            mean_IoU += iou_score\n",
    "            cumI += I\n",
    "            cumU += U\n",
    "        \n",
    "            # print(img_id, iou_score, flush=True)\n",
    "            for ii, thres in enumerate(thres_list):\n",
    "                if iou_score >= thres:\n",
    "                    count[ii]+=1\n",
    "            valid_count += 1\n",
    "\n",
    "        except Exception as e:\n",
    "            print(e, flush=True)\n",
    "            print('invalid output', img_id, pred_bbox, gt_bbox, flush=True)\n",
    "\n",
    "\n",
    "print('number of total/valid samples:', total_count, valid_count)\n",
    "for ii, thres in enumerate(thres_list):\n",
    "    print(f'Acc at iou_{thres}:', count[ii] / total_count * 100, flush=True)\n",
    "print(f'meanIoU: {mean_IoU/total_count * 100}, cumIoU: {cumI/cumU * 100}, rot_err: {np.mean(np.abs(rot_errs))}', flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Captioning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = os.path.join(output_dir, f'rsbench_cap{output_version}.json')\n",
    "# data_path = '../../outputs/GPT4_eval/gpt4_cap.json'\n",
    "\n",
    "gt_answers= []\n",
    "pred_answers = []\n",
    "with open(data_path, 'r') as file:\n",
    "    for line in file:\n",
    "        item = json.loads(line.strip())\n",
    "        img_id = item['image_id']\n",
    "        # if '09194_0000' in data_path:\n",
    "        #     print(item)\n",
    "\n",
    "        gt_ans = item['ground_truth'].strip().replace('\\n', ' ')\n",
    "        pred_ans = item['answer'].strip().replace('\\n', ' ')\n",
    "\n",
    "        if img_id is None or img_id=='\\n' or pred_ans is None or pred_ans=='\\n':\n",
    "            print('empty', img_id, pred_ans)\n",
    "            continue\n",
    "\n",
    "        gt_answers.append([img_id, gt_ans])\n",
    "        pred_answers.append([img_id, pred_ans])\n",
    "\n",
    "print('number of captions', len(gt_answers))\n",
    "\n",
    "np.savetxt('pred_cap.txt',  pred_answers, fmt='%s', delimiter='\\t')\n",
    "np.savetxt('gt_cap.txt',  gt_answers, fmt='%s', delimiter='\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cap_len = [len(ans[1].split()) for ans in pred_answers]\n",
    "print('avg len', np.mean(cap_len), np.std(cap_len))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# conda activate YOUR_PYTHON_ENV\n",
    "# cd geochat/eval/caption_eval/\n",
    "# python create_json_references.py -i ../gt_cap.txt -o ../gt_cap.json\n",
    "# python run_evaluations.py -i ../pred_cap.txt  -r ../gt_cap.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geochat",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
